/*******************************************************************************
 * Copyright (c) 2017-2019, org.smartboot. All rights reserved.
 * project name: smart-socket
 * file name: UdpBootstrap.java
 * Date: 2019-12-31
 * Author: sandao (zhengjunweimail@163.com)
 *
 ******************************************************************************/

package org.ricks.net.transport;

import org.ricks.common.conf.NetWorkConfig;
import org.ricks.ioc.Message;
import org.ricks.net.Server;
import org.ricks.net.buffer.BufferPage;
import org.ricks.net.buffer.BufferPagePool;
import org.ricks.net.buffer.VirtualBuffer;
import org.ricks.net.handler.MessageProcessor;
import org.ricks.net.handler.NetMonitor;
import org.ricks.net.StateMachineEnum;
import org.ricks.net.handler.Protocol;
import org.ricks.utils.ByteUtils;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.DatagramChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * UDP服务启动类
 *
 * @author 三刀
 * @version V1.0 , 2019/8/18
 */
public class UdpBootstrap extends Server {

    private final static int MAX_READ_TIMES = 16;

    /**
     * 缓存页
     */
    private final BufferPage bufferPage = new BufferPagePool(1024 * 1024, 1, true).allocateBufferPage();

    private Worker worker;

    private boolean running = true;

    public UdpBootstrap() {

    }

    public <Request> UdpBootstrap(Protocol<Request> protocol) {
        config.setProtocol(protocol);
    }

    public <Request> UdpBootstrap(Protocol<Request> protocol, MessageProcessor<Request> messageProcessor) {
        config.setProtocol(protocol);
        config.setProcessor(messageProcessor);
    }

    /**
     * 开启一个UDP通道，端口号随机
     *
     * @return UDP通道
     */
    public void start() {
        try {
            open(NetWorkConfig.getPort());
        } catch (IOException e) {
            System.err.println("====================================================");
            System.err.println("==  server started Failed port:" + NetWorkConfig.getPort() + "   ===");
            System.err.println("====================================================");
            e.printStackTrace();
        }
        Runtime.getRuntime().addShutdownHook(new Thread(() -> stop()));
    }

    /**
     * 开启一个UDP通道
     *
     * @param port 指定绑定端口号,为0则随机指定
     */
    public UdpChannel open(int port) throws IOException {
        return open(null, port);
    }

    /**
     * 开启一个UDP通道
     *
     * @param host 绑定本机地址
     * @param port 指定绑定端口号,为0则随机指定
     */
    public UdpChannel open(String host, int port) throws IOException {
        //启动线程服务
        if (worker == null) {
            initThreadServer();
            //启动Boss线程组
            worker = new Worker();
            new Thread(worker, "socket-0").start();
        }

        DatagramChannel channel = DatagramChannel.open();
        channel.configureBlocking(false);
        if (port > 0) {
            InetSocketAddress inetSocketAddress = host == null ? new InetSocketAddress(port) : new InetSocketAddress(host, port);
            channel.socket().bind(inetSocketAddress);
        }
        UdpChannel udpChannel = initChannel(channel);
        worker.addRegister(selector -> {
            try {
                SelectionKey selectionKey = channel.register(selector, SelectionKey.OP_READ);
                udpChannel.setSelectionKey(selectionKey);
                selectionKey.attach(udpChannel);
            } catch (ClosedChannelException e) {
                e.printStackTrace();
            }
        });
        return udpChannel;
    }

    protected UdpChannel initChannel(DatagramChannel channel) {
        return new UdpChannel(channel, worker, config, bufferPage);
    }

    protected KcpChannel createKcpChannel(DatagramChannel channel) {
        return new KcpChannel(channel, worker, config, bufferPage);
    }


    /**
     * Kcp 接受实际是经过UDP
     * @param readBuffer
     * @param channel
     * @throws IOException
     */
    private void doRead(VirtualBuffer readBuffer, UdpChannel channel) throws IOException {
        int count = MAX_READ_TIMES;
        while (count-- > 0) {
            //接收数据
            ByteBuffer buffer = readBuffer.buffer();
            buffer.clear();
            SocketAddress remote = channel.getChannel().receive(buffer);
            if (remote == null) {
                return;
            }
            byte[] bytes = new byte[buffer.position()];
            buffer.flip();
            buffer.get(bytes);
            UdpAioSession aioSession = channel.createAndCacheSession(remote); //创建Session
            readHandle(aioSession,buffer,bytes);
        }
    }

    protected void readHandle(UdpAioSession session,ByteBuffer buffer, byte[] bytes) {
        NetMonitor netMonitor = config.getMonitor();
        if (netMonitor != null) {
            netMonitor.beforeRead(session);
            netMonitor.afterRead(session, buffer.remaining());
        }
        ByteBuffer data = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
        receiveHandle(session,data); //KCP下放处理
    }

    protected void receiveHandle(UdpAioSession session,ByteBuffer data) {
        session.setLastRecieveTime(System.currentTimeMillis());
        Message request = null;
        //解码
        try {
            request = (Message) config.getProtocol().decode(data, session);
        } catch (Exception e) {
            config.getProcessor().stateEvent(session, StateMachineEnum.DECODE_EXCEPTION, e);
            e.printStackTrace();
        }
        dispatch(session,request,session.getRemoteAddress());
    }

    public void shutdown() {
        running = false;
        worker.selector.wakeup();
    }

    /**
     * 设置读缓存区大小
     *
     * @param size 单位：byte
     */
    public final UdpBootstrap setReadBufferSize(int size) {
        this.config.setReadBufferSize(size);
        return this;
    }


    /**
     * 设置线程大小
     *
     * @param num 线程数
     */
    public final UdpBootstrap setThreadNum(int num) {
        this.config.setThreadNum(num);
        return this;
    }


    /**
     * 是否启用控制台Banner打印
     *
     * @param bannerEnabled true:启用，false:禁用
     * @return 当前AioQuickServer对象
     */
    public final UdpBootstrap setBannerEnabled(boolean bannerEnabled) {
        config.setBannerEnabled(bannerEnabled);
        return this;
    }

    class Worker implements Runnable {
        /**
         * 当前Worker绑定的Selector
         */
        private final Selector selector;

        /**
         * 待注册的事件
         */
        private final ConcurrentLinkedQueue<Consumer<Selector>> registers = new ConcurrentLinkedQueue<>();

        Worker() throws IOException {
            this.selector = Selector.open();
        }

        /**
         * 注册事件
         */
        final void addRegister(Consumer<Selector> register) {
            registers.offer(register);
            selector.wakeup();
        }

        @Override
        public final void run() {
            // 优先获取SelectionKey,若无关注事件触发则阻塞在selector.select(),减少select被调用次数
            Set<SelectionKey> keySet = selector.selectedKeys();
            //读缓冲区
            VirtualBuffer readBuffer = bufferPage.allocate(config.getReadBufferSize());
            try {
                while (running) {
                    Consumer<Selector> register;
                    while ((register = registers.poll()) != null) {
                        register.accept(selector);
                    }
                    if (keySet.isEmpty() && selector.select() == 0) {
                        continue;
                    }
                    Iterator<SelectionKey> keyIterator = keySet.iterator();
                    // 执行本次已触发待处理的事件
                    while (keyIterator.hasNext()) {
                        SelectionKey key = keyIterator.next();
                        keyIterator.remove();
                        UdpChannel udpChannel = (UdpChannel) key.attachment();
                        if (!key.isValid()) {
                            udpChannel.close();
                            continue;
                        }
                        if (key.isReadable()) {
                            doRead(readBuffer, udpChannel);
                        }
                        if (key.isWritable()) {
                            udpChannel.doWrite();
                        }
                    }
                }
            } catch (Exception e) {
                stop();
                e.printStackTrace();
            } finally {
                //读缓冲区内存回收
                readBuffer.clean();
            }
        }
    }
}

